import sys
import json
from IPython.display import Markdown, display
from llama_index.llms.ollama import Ollama
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.embeddings.fastembed import FastEmbedEmbedding
from llama_index.core import SQLDatabase
from llama_index.core import Settings
from llama_index.core.query_engine import NLSQLTableQueryEngine
from sqlalchemy import create_engine
import openai
import logging
from llama_index.llms.openai import OpenAI
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
#pip install beautifulsoup4 lxml
import uuid
from bs4 import BeautifulSoup


# from actions.api.mysql_tab import MysqlTab
logger = logging.getLogger(__name__)

class RagUI:
    """ RagUI """

    def __init__(self ):
        logger.info(' RagUI -init')
        openai.api_key = ''
        openai.base_url = ''
        
        # self.llm = OpenAI(temperature=0.19, model="gpt-4-turbo-preview")
        self.llm = Ollama(model="codeqwen:7b-chat", base_url="http://localhost:11434",request_timeout=60000.0)
        self.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-zh-v1.5")        
        Settings.llm = self.llm
        Settings.embed_model = self.embed_model
        
    def rag(self,query_str):
        prompt=f'''
        按要求生成html5页面,语言是中文。        
        要求是:{query_str}
        请只回复答案，不要写解释。
        '''
        # html_str = self.llm.complete(prompt)
        # print(html_str.text)
        # html = html_str.text
        # 假设这是你的HTML字符串
        html_string = """
<html>
 <head>
  <script src="https://cdn.tailwindcss.com">
  </script>
  <link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css" rel="stylesheet"/>
 </head>
 <body class="flex flex-col items-center justify-center min-h-screen bg-white">
  <div class="flex flex-col items-center">
   <img alt="Baidu logo with red and blue text" class="mb-6" height="100" src="https://oaidalleapiprodscus.blob.core.windows.net/private/org-CA2iYIwLQOEhWnyU3Nh1e7Va/user-OiBpZoAJVXnhd2tyYm2U0O5W/img-XsiRykhtSA8boD2yy9mnjOXV.png?st=2024-05-15T07%3A45%3A16Z&amp;se=2024-05-15T09%3A45%3A16Z&amp;sp=r&amp;sv=2021-08-06&amp;sr=b&amp;rscd=inline&amp;rsct=image/png&amp;skoid=6aaadede-4fb3-4698-a8f6-684d7786b067&amp;sktid=a48cca56-e6da-484e-a814-9c849652bcb3&amp;skt=2024-05-14T15%3A23%3A30Z&amp;ske=2024-05-15T15%3A23%3A30Z&amp;sks=b&amp;skv=2021-08-06&amp;sig=YyfSi7eGZkXIc07hRA90xrsFndA7T5hbutg6fQmAr9E%3D" width="200"/>
   <div class="flex items-center border-2 border-blue-500 rounded-full overflow-hidden">
    <input class="w-96 h-12 px-4 outline-none" placeholder="" type="text"/>
    <button class="flex items-center justify-center w-12 h-12 bg-blue-500 text-white">
     <i class="fas fa-camera">
     </i>
    </button>
    <button class="px-4 h-12 bg-blue-500 text-white text-lg">
     百度一下
    </button>
   </div>
  </div>
 </body>
</html>
"""
        div = ""
        # 使用BeautifulSoup加载HTML字符串
        soup = BeautifulSoup(html_string, 'html.parser')
        # 找到<body>标签
        body = soup.find('body')
        # 遍历<body>标签的所有子对象
        for child in body.children:
            # 检查child是否是一个Tag对象
            if child.name is  None:
                continue
            # 打印子对象的完整内部HTML
            print(f"Tag: {child.name}, InnerHTML: {str(child)}")
            if child.name.lower() == 'div':
                div += ragUI.readUIDoc(str(child))
            if child.name.lower() == 'input':
                div += ragUI.readUIDoc(str(child))
            if child.name.lower() == 'button':
                div += ragUI.readUIDoc(str(child))
        new_uuid = str(uuid.uuid4())
        page ='{"id":"{'+new_uuid+'}","style":{"padding-left":"20px","padding-right":"20px","padding-bottom":"20px"},"FolderId":null,"ParentPageId":null,"children":['+div+'],"containerLeavel":0,"name":"Page","label":"page","type":"page","code":"apj_page","moduleId":"811C0F4CEA8E"}'
        return page
        
    def readUIDoc(self,query_str):        
        db = chromadb.PersistentClient(path=r"E:\100rag\chroma")
        chroma_collection = db.get_or_create_collection("doc-ui")
        vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
        index = VectorStoreIndex.from_vector_store(
            vector_store,
            embed_model=self.embed_model,
        )
        
        # query
        prompt=f'''实现功能如下:
1:字符串替换功能，不输出代码，输出替换后的结果.
2:生成随机生成8位字符串替换字符串里面uuid.
3:请只回复答案，不要写解释.
请替换下面字符串:{query_str}
        '''
        # query_engine = index.as_query_engine(similarity_top_k=2, llm=gpt35_llm)
        doc_engine = index.as_query_engine(similarity_top_k=6)
        rag_doc = doc_engine.query(prompt)
        print('#########################')
        # print(rag_doc)
        if hasattr(rag_doc, 'response'):
            return rag_doc.response
        return '无法识别'
        

ragUI = RagUI()
ret = ragUI.rag("生成登录页面")
print(ret)
#ret = ragUI.readUIDoc('<button class="px-4 py-2 bg-blue-500 text-white flex items-center"><i class="fas fa-camera mr-2"></i>提交</button>')
#print(ret)
#ret = ragUI.readUIDoc('<input type="text" id="cargo" name="cargo" placeholder="请输入商品">')
#print(ret)
#ret = ragUI.readUIDoc('<div class="flex flex-col items-center"></div>')
#print(ret)
html='''
<div class="flex items-center border-2 border-blue-500 rounded-full overflow-hidden">
	<input class="px-4 py-2 w-96 outline-none" placeholder="" type="text"/>
	<button class="px-4 py-2 bg-blue-500 text-white flex items-center">
	<i class="fas fa-camera mr-2"></i>
	百度一下
	</button>
</div>
'''
#ret = ragUI.readUIDoc(html)
#print(ret)